#include <RcppArmadillo.h>

/* bapvar_mean() takes in a matrix of lagged outcome observations,
 * a matrix of lag coefficients, and the matrix \( \tilde{\lambda} \)
 * as defined in Appendix A.1 of Brandt and Sandler (2012)
 * (discussed in utils.cpp), and returns the BaP-VAR mean
 * conditional on beta, A, and D
 */
arma::mat bapvar_mean(const arma::mat& ylags, const arma::mat& A,
                      const arma::mat& lambda_tilde_mat) {
    arma::mat ylagsA = ylags * A;
    return ylagsA + lambda_tilde_mat;
}

/* bapvar_var() takes in a matrix of lagged outcome observations,
 * a matrix of lag coefficients, the matrix \( \tilde{\lambda} \)
 * as defined in Appendix A.1 of Brandt and Sandler (2012)
 * (discussed in utils.cpp), and a matrix D giving the covariance for b,
 * and returns the BaP-VAR variance conditional on beta, A, and D
 */
arma::mat bapvar_var(const arma::mat& ylags, const arma::mat& A,
                     const arma::mat& lambda_tilde_mat, const arma::mat& D) {
    const int m = D.n_cols;
    const int n = ylags.n_rows;
    // This differs from equation A4 in Brandt and Sandler (2012)
    // because the way I store A is the transpose of how they do
    arma::mat A_cov = A.t() * arma::cov(ylags) * A;
    arma::mat D_part = arma::exp(D) - arma::ones<arma::mat>(m, m);
    arma::mat Sigma = arma::zeros<arma::mat>(m, m);
    arma::mat Lambda_tilde(m, m);
    for ( arma::uword i = 0; i < n; ++i ) {
        Lambda_tilde = arma::diagmat(lambda_tilde_mat.row(i));
        Sigma += Lambda_tilde + Lambda_tilde * D_part * Lambda_tilde;
    }
    arma::mat result = A_cov + Sigma;
    result = result / n;
    return result;
}

